"""分析和可视化 trajectories.npz 文件"""
import numpy as np
import argparse


def analyze_trajectories(npz_path):
    """分析轨迹数据"""
    print(f"\n{'='*60}")
    print(f"加载轨迹数据: {npz_path}")
    print(f"{'='*60}")

    # 加载数据
    data = np.load(npz_path, allow_pickle=True)

    # 查看包含的数据
    print("\n📦 数据内容:")
    for key in data.files:
        print(f"  - {key}: {type(data[key])}")

    # 获取基本统计信息
    episode_rewards = data['episode_rewards']
    episode_lengths = data['episode_lengths']
    trajectories = data['trajectories']

    print(f"\n📊 基本统计:")
    print(f"  总 Episodes: {len(episode_rewards)}")
    print(f"  平均奖励: {np.mean(episode_rewards):.2f} ± {np.std(episode_rewards):.2f}")
    print(f"  最高奖励: {np.max(episode_rewards):.2f}")
    print(f"  最低奖励: {np.min(episode_rewards):.2f}")
    print(f"  平均长度: {np.mean(episode_lengths):.2f} 步")

    # 分析每个 episode
    print(f"\n📝 每个 Episode 的详细信息:")
    print(f"{'Episode':<10} {'奖励':<10} {'长度':<10} {'成功':<10}")
    print("-" * 40)

    for i, (reward, length) in enumerate(zip(episode_rewards, episode_lengths)):
        success = "✓" if reward > 0 else "✗"
        print(f"{i+1:<10} {reward:<10.2f} {int(length):<10} {success:<10}")

    # 分析最好的 episode
    best_idx = np.argmax(episode_rewards)
    best_traj = trajectories[best_idx]

    print(f"\n🏆 最佳 Episode (Episode {best_idx + 1}):")
    print(f"  总奖励: {episode_rewards[best_idx]:.2f}")
    print(f"  总步数: {int(episode_lengths[best_idx])}")

    # 分析轨迹详情
    if isinstance(best_traj, dict):
        print(f"\n  轨迹数据:")
        print(f"    - 观察序列长度: {len(best_traj.get('observations', []))}")
        print(f"    - 动作序列长度: {len(best_traj.get('actions', []))}")
        print(f"    - 奖励序列长度: {len(best_traj.get('rewards', []))}")

        # 显示动作序列
        actions = best_traj.get('actions', [])
        if actions:
            print(f"\n  动作序列（前20步）:")
            for step, action in enumerate(actions[:20]):
                print(f"    步骤 {step+1}: {action}")

        # 显示奖励序列
        rewards = best_traj.get('rewards', [])
        if rewards:
            print(f"\n  奖励分布:")
            reward_steps = [(i+1, r) for i, r in enumerate(rewards) if r > 0]
            if reward_steps:
                print(f"    获得奖励的步骤:")
                for step, reward in reward_steps:
                    print(f"      步骤 {step}: +{reward}")
            else:
                print(f"    所有步骤奖励均为 0")

            print(f"    累计奖励: {sum(rewards):.2f}")

    # 动作分布统计
    print(f"\n📈 全局动作分布:")
    all_actions = []
    for traj in trajectories:
        if isinstance(traj, dict) and 'actions' in traj:
            for action_pair in traj['actions']:
                if isinstance(action_pair, (list, np.ndarray)):
                    all_actions.extend(action_pair)
                else:
                    all_actions.append(action_pair)

    if all_actions:
        unique_actions, counts = np.unique(all_actions, return_counts=True)
        print(f"  动作空间: {unique_actions}")
        print(f"  动作统计:")
        for action, count in zip(unique_actions, counts):
            percentage = count / len(all_actions) * 100
            print(f"    动作 {action}: {count:>5} 次 ({percentage:>5.1f}%)")

    # 奖励随时间变化
    print(f"\n📉 Episode 奖励趋势:")
    for i in range(0, len(episode_rewards), max(1, len(episode_rewards)//10)):
        print(f"  Episode {i+1:>3}: {episode_rewards[i]:.2f}")

    return data


def export_best_trajectory(npz_path, output_path="best_trajectory.txt"):
    """导出最佳轨迹到文本文件"""
    data = np.load(npz_path, allow_pickle=True)

    episode_rewards = data['episode_rewards']
    episode_lengths = data['episode_lengths']
    trajectories = data['trajectories']

    best_idx = np.argmax(episode_rewards)
    best_traj = trajectories[best_idx]

    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("="*60 + "\n")
        f.write(f"最佳 Episode 详细轨迹 (Episode {best_idx + 1})\n")
        f.write("="*60 + "\n\n")
        f.write(f"总奖励: {episode_rewards[best_idx]:.2f}\n")
        f.write(f"总步数: {int(episode_lengths[best_idx])}\n\n")

        if isinstance(best_traj, dict):
            actions = best_traj.get('actions', [])
            rewards = best_traj.get('rewards', [])
            observations = best_traj.get('observations', [])

            f.write("-"*60 + "\n")
            f.write(f"{'步骤':<6} {'动作(Agent0,Agent1)':<25} {'即时奖励':<12} {'累计奖励':<12}\n")
            f.write("-"*60 + "\n")

            cumulative_reward = 0
            for step, (action, reward) in enumerate(zip(actions, rewards)):
                cumulative_reward += reward
                action_str = str(action) if isinstance(action, (list, tuple)) else f"[{action}]"
                f.write(f"{step+1:<6} {action_str:<25} {reward:<12.2f} {cumulative_reward:<12.2f}\n")

                if reward > 0:
                    f.write(f"       >>> 获得奖励！\n")

    print(f"\n✓ 最佳轨迹已导出到: {output_path}")


def compare_episodes(npz_path, episode_ids=None):
    """比较多个 episode"""
    data = np.load(npz_path, allow_pickle=True)

    episode_rewards = data['episode_rewards']
    episode_lengths = data['episode_lengths']
    trajectories = data['trajectories']

    if episode_ids is None:
        # 选择最好和最差的几个
        sorted_indices = np.argsort(episode_rewards)
        episode_ids = [sorted_indices[0], sorted_indices[-1]]  # 最差和最好

    print(f"\n📊 Episode 对比:")
    print(f"{'Metric':<20} " + "  ".join([f"Episode {i+1:>3}" for i in episode_ids]))
    print("-" * (20 + 15 * len(episode_ids)))

    # 奖励
    rewards_str = "  ".join([f"{episode_rewards[i]:>10.2f}" for i in episode_ids])
    print(f"{'总奖励':<20} {rewards_str}")

    # 长度
    lengths_str = "  ".join([f"{int(episode_lengths[i]):>10}" for i in episode_ids])
    print(f"{'Episode长度':<20} {lengths_str}")

    # 动作多样性
    for idx in episode_ids:
        traj = trajectories[idx]
        if isinstance(traj, dict) and 'actions' in traj:
            actions = traj['actions']
            unique_actions = len(set(str(a) for a in actions))
            print(f"Episode {idx+1} 动作多样性: {unique_actions} 种不同动作")


def main():
    parser = argparse.ArgumentParser(description='分析轨迹数据')
    parser.add_argument('npz_file', type=str, help='trajectories.npz 文件路径')
    parser.add_argument('--export', action='store_true', help='导出最佳轨迹到文本文件')
    parser.add_argument('--compare', type=int, nargs='+', help='比较指定的 episodes (从1开始计数)')

    args = parser.parse_args()

    # 分析轨迹
    data = analyze_trajectories(args.npz_file)

    # 导出最佳轨迹
    if args.export:
        output_path = args.npz_file.replace('.npz', '_best.txt')
        export_best_trajectory(args.npz_file, output_path)

    # 比较 episodes
    if args.compare:
        episode_ids = [i-1 for i in args.compare]  # 转换为0-based索引
        compare_episodes(args.npz_file, episode_ids)

    print(f"\n{'='*60}")
    print("分析完成！")
    print(f"{'='*60}\n")


if __name__ == "__main__":
    main()
